from dataclasses import dataclass
import os
pdg_generator='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/pdg_generation'
classifier='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/classification/classifier.py'
learner='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/classification/learner.py'

@dataclass
class JStapMeta:
    train_dir:str=""
    validate_dir:str=''
    test_dir:str=''
    analysis_dir:str=''
    is_origin:bool=False
    
    def model_path(self,model_name='ast_ngrams'):
        if self.is_origin:
            model_name='ngrams_ast'
        return f'{self.analysis_dir}/Model/{model_name}'

    @property
    def train_good(self):
        
        if self.is_origin:
            return f'{self.train_dir}/Good-train'

        return f'{self.train_dir}/good'
    
    @property
    def train_bad(self):
        if self.is_origin:
            return f'{self.train_dir}/Bad-train'
        
        return f'{self.train_dir}/bad'
    
    @property
    def validate_good(self):
        if self.is_origin:
            return f'{self.validate_dir}/Good-validate'
        
        return f'{self.validate_dir}/good'
    
    @property
    def validate_bad(self):
        if self.is_origin:
            return f'{self.validate_dir}/Bad-validate'
        
        return f'{self.validate_dir}/bad'
    
    @property
    def test_good(self):
        if self.is_origin:
            return f'{self.test_dir}/Good-test'
        
        return f'{self.test_dir}/good'
    
    @property
    def test_bad(self):
        if self.is_origin:
            return f'{self.test_dir}/Bad-test'
        
        return f'{self.test_dir}/bad'
    
    @staticmethod
    def auto(root_dir:str):
        analysis_dir=os.path.join(root_dir,'Anaysis')
        if os.path.exists(os.path.join(root_dir,'train')):
            train_dir=os.path.join(root_dir,'train')
            validate_dir=os.path.join(root_dir,'validate')
            test_dir=os.path.join(root_dir,'test')
            
        else:
            train_dir=os.path.join(root_dir,'jstap_train')
            validate_dir=os.path.join(root_dir,'jstap_validate')    
            test_dir=os.path.join(root_dir,'jstap_test')
        
        return JStapMeta(train_dir,validate_dir,test_dir,analysis_dir)

    @staticmethod
    def origin(root_dir:str):
        analysis_dir=os.path.join(root_dir,'Anaysis')
        samples_dir=os.path.join(root_dir,'samples')
        return JStapMeta(train_dir=samples_dir,validate_dir=samples_dir,test_dir=samples_dir,analysis_dir=analysis_dir,is_origin=True)

def build_pdg(dir,rebuild=False):
    if os.path.exists(os.path.join(dir,'Analysis')) and not rebuild:
        return
    #cmd=f'python -c "import sys; sys.path.insert(0, "{pdg_generator}"); from pdgs_generation import *; store_pdg_folder("{dir}")"'
    cmd=f"python -c \"import sys; sys.path.insert(0, '{pdg_generator}'); from pdgs_generation import *; store_pdg_folder('{dir}')\""
    return os.system(cmd)

def classify(model_path,good_dir,bad_dir,anaysis_path):
    level='ast'
    feature_choice='ngrams'

    cmd=f"python {classifier} --d {good_dir}/Analysis/PDG  {bad_dir}/Analysis/PDG --l benign malicious --level {level} --features {feature_choice} --m {model_path} --analysis_path {anaysis_path}"

    return os.system(cmd)

def classify_origin(good_dir,bad_dir):
    model_path='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/Analysis/Model/ngrams_ast'
    analysis_dir='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/Analysis'

    return classify(model_path,good_dir,bad_dir,analysis_dir)

def train(good_train_dir,bad_train_dir,good_validate_dir,bad_validate_dir,analysis_path):
    level='ast'
    feature_choice='ngrams'
    model_dir=f'{analysis_path}/Model'
    model_name=f'{level}_{feature_choice}'
    cmd=f"python {learner} --d {good_train_dir}/Analysis/PDG  {bad_train_dir}/Analysis/PDG --l benign malicious --vd {good_validate_dir}/Analysis/PDG  {bad_validate_dir}/Analysis/PDG --vl benign malicious --level {level}  --features {feature_choice} --mn {model_name} --md {model_dir} --analysis_path {analysis_path}"

    return f'{model_dir}/{model_name}',os.system(cmd)
    #"--vd $samples_path'Good-validate'$pdg $samples_path'Bad-validate'$pdg --vl benign malicious --level $level --features $features_choice --mn $model_name --md $model_dir"



def mix_dataset(dataset_pairs,out_dir):
    os.makedirs(out_dir,exist_ok=True)
    os.makedirs(os.path.join(out_dir,'Analysis/PDG'),exist_ok=True)
    for dataset_dir,label in dataset_pairs:
        files=os.listdir(dataset_dir)
        files=[os.path.join(dataset_dir,file) for file in files if file.endswith('.js')]
        
        analysis_dir=os.path.join(dataset_dir,'Analysis/PDG')
        pdg_files=[os.path.join(analysis_dir,os.path.basename(file)[:-3]) for file in files]
        
        
        for file,pdg_file in zip(files,pdg_files):
            #put the label to the file name
            out_file=os.path.join(out_dir,f'{label}_{os.path.basename(file)}')
            out_pdg_file=os.path.join(out_dir,'Analysis/PDG',f'{label}_{os.path.basename(pdg_file)}')
            #copy file to out_dir
            #print(f'{file} {out_file}')
            #print((f'{pdg_file} {out_pdg_file}'))
            os.system(f'cp {file} {out_file}')
            if os.path.exists(pdg_file):
                os.system(f'cp {pdg_file} {out_pdg_file}')

    @dataclass
    class JStapMeta:
        train_dir:str=""
        validate_dir:str=''
        test_dir:str=''
        analysis_dir:str=''
        
        def model_path(self,model_name='ast_ngrams'):
            return f'{self.analysis_dir}/Model/{model_name}'

        @property
        def train_good(self):
            return f'{self.train_dir}/good'
        
        @property
        def train_bad(self):
            return f'{self.train_dir}/bad'
        
        @property
        def validate_good(self):
            return f'{self.validate_dir}/good'
        
        @property
        def validate_bad(self):
            return f'{self.validate_dir}/bad'
        
        @property
        def test_good(self):
            return f'{self.test_dir}/good'
        
        @property
        def test_bad(self):
            return f'{self.test_dir}/bad'
        
        @staticmethod
        def auto(root_dir:str):
            analysis_dir=os.path.join(root_dir,'Anaysis')
            if os.path.exists(os.path.join(root_dir,'train')):
                train_dir=os.path.join(root_dir,'train')
                validate_dir=os.path.join(root_dir,'validate')
                test_dir=os.path.join(root_dir,'test')
                
            elif os.path.exists(os.path.join(root_dir,'jstap_train')):
                train_dir=os.path.join(root_dir,'jstap_train')
                validate_dir=os.path.join(root_dir,'jstap_validate')    
                test_dir=os.path.join(root_dir,'jstap_test')
            else:
                return None
            
            return JStapMeta(train_dir,validate_dir,test_dir,analysis_dir)

def train_meta(meta:JStapMeta):
        return train(meta.train_good,meta.train_bad,meta.validate_good,meta.validate_bad,meta.analysis_dir)
    
def classify_meta(model_path,meta:JStapMeta,analysis_dir=None):
    if analysis_dir is None:
        analysis_dir=meta.analysis_dir
    return classify(model_path,meta.test_good,meta.test_bad,analysis_dir)

def mix_dataset_meta(meta0:JStapMeta,origin_label,meta1:JStapMeta,new_label,out_meta:JStapMeta):
    pairs_0=[(meta0.train_good,origin_label),(meta1.train_good,new_label)]
    pairs_1=[(meta0.validate_good,origin_label),(meta1.validate_good,new_label)]
    pairs_2=[(meta0.test_good,origin_label),(meta1.test_good,new_label)]
    pairs_3=[(meta0.train_bad,origin_label),(meta1.train_bad,new_label)]
    pairs_4=[(meta0.validate_bad,origin_label),(meta1.validate_bad,new_label)]
    pairs_5=[(meta0.test_bad,origin_label),(meta1.test_bad,new_label)]

    mix_dataset(pairs_0,out_meta.train_good)
    mix_dataset(pairs_1,out_meta.validate_good)
    mix_dataset(pairs_2,out_meta.test_good)
    mix_dataset(pairs_3,out_meta.train_bad)
    mix_dataset(pairs_4,out_meta.validate_bad)
    mix_dataset(pairs_5,out_meta.test_bad)

if __name__=='__main__':
    

        
    

    model_path='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/Analysis/Model/ngrams_ast'


    obfuscate_medium_dataset=JStapMeta.auto('/home/passwd123/XiaoweiGuo/vscode/code-security/data/js/obfucscate/medium')
    obfuscate_light_dataset=JStapMeta.auto('/home/passwd123/XiaoweiGuo/vscode/code-security/data/js/obfucscate/light')
    obfuscate_mix_dataset=JStapMeta.auto('/home/passwd123/XiaoweiGuo/vscode/code-security/data/js/obfucscate/mix/md')
    obuscate_mix_ld_dataset=JStapMeta.auto('/home/passwd123/XiaoweiGuo/vscode/code-security/data/js/obfucscate/mix/ld')

    good_train_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Good-train'
    bad_train_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Bad-train'

    bad_validate_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Bad-validate'
    good_validate_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Good-validate'

    good_test_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Good-test'
    bad_test_dir_origin='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/samples/Bad-test'

    #out_dir='/home/passwd123/XiaoweiGuo/vscode/code-security/exp_result/JStap/Obfuscation/mix/'

    out_dir='/home/passwd123/XiaoweiGuo/vscode/code-security/third_party/JStap/Analysis'


    #out_dir='/home/passwd123/XiaoweiGuo/vscode/code-security/data/js/obfucscate/mix/ld'

    pairs_good=[(obfuscate_light_dataset.test_good,'light'),(good_test_dir_origin,'origin')]
    pairs_bad=[(obfuscate_light_dataset.test_bad,'light'),(bad_test_dir_origin,'origin')]

    #mix_dataset(pairs_good,out_dir+'/test/good')
    #mix_dataset(pairs_bad,out_dir+'/test/bad')
    # build_pdg(good_validate_dir)
    # build_pdg(bad_validate_dir)

    # build_pdg(good_dir)
    # build_pdg(bad_dir)

    #model_path,stat=train_meta(obuscate_mix_ld_dataset)

    classify_meta(obfuscate_medium_dataset.model_path(),obfuscate_light_dataset,obfuscate_medium_dataset.analysis_dir)

    #classify(model_path,good_test_dir_origin,bad_test_dir_origin,out_dir)